import pathlib

import torch.utils.data

import torch_plus
from lotka_volterra_problem import LotkaVolterraProblem


class LotkaVolterraDataset:
    def __init__(self,
                 problem: LotkaVolterraProblem,
                 t1_batch: torch.Tensor,
                 t2_batch: torch.Tensor):
        self._problem = problem
        self._t1 = torch_plus.clone_detached(t1_batch)
        self._t2 = torch_plus.clone_detached(t2_batch)

        self._t_min = torch.min(torch.min(self._t1), torch.min(self._t2))
        self._t_max = torch.max(torch.max(self._t1), torch.max(self._t2))

        x1_batch = [problem.x1(float(t)) for t in self._t1]
        x1_batch = torch.tensor(x1_batch)
        self._x1 = x1_batch

        x2_batch = [problem.x2(float(t)) for t in self._t2]
        x2_batch = torch.tensor(x2_batch)
        self._x2 = x2_batch

    def problem(self) -> LotkaVolterraProblem:
        return self._problem

    def t1(self):
        return self._t1

    def t2(self):
        return self._t2

    def t_min(self):
        return self._t_min

    def t_max(self):
        return self._t_max

    def x1(self):
        return self._x1

    def x2(self):
        return self._x2

    def save(self, path: pathlib.Path):
        torch.save({
            "t1": self._t1,
            "t2": self._t2,
            "problem": {
                "a1": self._problem.a1(),
                "b1": self._problem.b1(),
                "a2": self._problem.a2(),
                "b2": self._problem.b2(),
                "x1_0": self._problem.x1_0(),
                "x2_0": self._problem.x2_0(),
                "t_step": self._problem.t_step(),
            }
        }, path)

    @staticmethod
    def load(path: pathlib.Path) -> 'LotkaVolterraDataset':
        data = torch.load(path)
        problem = data["problem"]
        return LotkaVolterraDataset(
            LotkaVolterraProblem(
                problem["a1"],
                problem["b1"],
                problem["a2"],
                problem["b2"],
                problem["x1_0"],
                problem["x2_0"],
                problem["t_step"]
            ),
            data["t1"],
            data["t2"]
        )
